from sqlalchemy import Column, Integer, Index, String
from sqlalchemy.orm.session import Session

from airflow.models.base import Base
from airflow.utils.session import provide_session


class Folder(Base):

    __tablename__ = "idata_folder"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project_id = Column(Integer)
    name = Column(String(255))
    category = Column(Integer)
    parent_id = Column(Integer)
    level = Column(Integer)

    __table_args__ = (
        Index('parent_id', parent_id),
        Index('project_id', project_id),
        Index('project_id_name', project_id, name),
    )

    @staticmethod
    @provide_session
    def get_folder(folder_id, session: Session = None):
        """
        Get the Task with specific task id.

        :param task_id: The task id of the Task to get.
        :param session: SQLAlchemy ORM Session
        :return: the Task object
        """
        return session.query(Folder).filter(Folder.id == folder_id).first()

    @staticmethod
    @provide_session
    def query_children(folder_id, session: Session = None):
        return session.query(Folder).filter(Folder.parent_id == folder_id).order_by(Folder.name).all()

    @staticmethod
    @provide_session
    def query(project_id, category=None, level=None, session: Session = None):
        qry = session.query(Folder).filter(Folder.project_id == project_id)
        if category:
            qry.filter(Folder.category == category)
        if level:
            qry.filter(Folder.level == level)
        return qry.all()

    @provide_session
    def insert(self, session: Session = None):
        session.add(self)

    @provide_session
    def update(self, session: Session = None):
        session.merge(self)

    def to_json(self):
        data = {'id': self.id, 'project_id': self.project_id, 'name': self.name, 'parent_id': self.parent_id,
                'level': self.level, 'category': self.category}
        return data
